{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 文件名称   ：roxml_to_dota.py\n",
    "# 功能描述   ：把rolabelimg标注的xml文件转换成dota能识别的xml文件，\n",
    "#             再转换成dota格式的txt文件\n",
    "#            把旋转框 cx,cy,w,h,angle，或者矩形框cx,cy,w,h,转换成四点坐标x1,y1,x2,y2,x3,y3,x4,y4\n",
    "import os\n",
    "import xml.etree.ElementTree as ET\n",
    "import math\n",
    " \n",
    "cls_list = ['site', 'wall']  #修改为自己的标签\n",
    " \n",
    " \n",
    "def edit_xml(xml_file, dotaxml_file):\n",
    "    \"\"\"\n",
    "    修改xml文件\n",
    "    :param xml_file:xml文件的路径\n",
    "    :return:\n",
    "    \"\"\"\n",
    " \n",
    "    # dxml_file = open(xml_file,encoding='gbk')\n",
    "    # tree = ET.parse(dxml_file).getroot()\n",
    " \n",
    "    tree = ET.parse(xml_file)\n",
    "    objs = tree.findall('object')\n",
    "    for ix, obj in enumerate(objs):\n",
    "        x0 = ET.Element(\"x0\")  # 创建节点\n",
    "        y0 = ET.Element(\"y0\")\n",
    "        x1 = ET.Element(\"x1\")\n",
    "        y1 = ET.Element(\"y1\")\n",
    "        x2 = ET.Element(\"x2\")\n",
    "        y2 = ET.Element(\"y2\")\n",
    "        x3 = ET.Element(\"x3\")\n",
    "        y3 = ET.Element(\"y3\")\n",
    "        # obj_type = obj.find('bndbox')\n",
    "        # type = obj_type.text\n",
    "        # print(xml_file)\n",
    " \n",
    "        if (obj.find('robndbox') == None):\n",
    "            obj_bnd = obj.find('bndbox')\n",
    "            obj_xmin = obj_bnd.find('xmin')\n",
    "            obj_ymin = obj_bnd.find('ymin')\n",
    "            obj_xmax = obj_bnd.find('xmax')\n",
    "            obj_ymax = obj_bnd.find('ymax')\n",
    "            # 以防有负值坐标\n",
    "            xmin = max(float(obj_xmin.text), 0)\n",
    "            ymin = max(float(obj_ymin.text), 0)\n",
    "            xmax = max(float(obj_xmax.text), 0)\n",
    "            ymax = max(float(obj_ymax.text), 0)\n",
    "            obj_bnd.remove(obj_xmin)  # 删除节点\n",
    "            obj_bnd.remove(obj_ymin)\n",
    "            obj_bnd.remove(obj_xmax)\n",
    "            obj_bnd.remove(obj_ymax)\n",
    "            x0.text = str(xmin)\n",
    "            y0.text = str(ymax)\n",
    "            x1.text = str(xmax)\n",
    "            y1.text = str(ymax)\n",
    "            x2.text = str(xmax)\n",
    "            y2.text = str(ymin)\n",
    "            x3.text = str(xmin)\n",
    "            y3.text = str(ymin)\n",
    "        else:\n",
    "            obj_bnd = obj.find('robndbox')\n",
    "            obj_bnd.tag = 'bndbox'  # 修改节点名\n",
    "            obj_cx = obj_bnd.find('cx')\n",
    "            obj_cy = obj_bnd.find('cy')\n",
    "            obj_w = obj_bnd.find('w')\n",
    "            obj_h = obj_bnd.find('h')\n",
    "            obj_angle = obj_bnd.find('angle')\n",
    "            cx = float(obj_cx.text)\n",
    "            cy = float(obj_cy.text)\n",
    "            w = float(obj_w.text)\n",
    "            h = float(obj_h.text)\n",
    "            angle = float(obj_angle.text)\n",
    "            obj_bnd.remove(obj_cx)  # 删除节点\n",
    "            obj_bnd.remove(obj_cy)\n",
    "            obj_bnd.remove(obj_w)\n",
    "            obj_bnd.remove(obj_h)\n",
    "            obj_bnd.remove(obj_angle)\n",
    " \n",
    "            x0.text, y0.text = rotatePoint(cx, cy, cx - w / 2, cy - h / 2, -angle)\n",
    "            x1.text, y1.text = rotatePoint(cx, cy, cx + w / 2, cy - h / 2, -angle)\n",
    "            x2.text, y2.text = rotatePoint(cx, cy, cx + w / 2, cy + h / 2, -angle)\n",
    "            x3.text, y3.text = rotatePoint(cx, cy, cx - w / 2, cy + h / 2, -angle)\n",
    " \n",
    "        # obj.remove(obj_type)  # 删除节点\n",
    "        obj_bnd.append(x0)  # 新增节点\n",
    "        obj_bnd.append(y0)\n",
    "        obj_bnd.append(x1)\n",
    "        obj_bnd.append(y1)\n",
    "        obj_bnd.append(x2)\n",
    "        obj_bnd.append(y2)\n",
    "        obj_bnd.append(x3)\n",
    "        obj_bnd.append(y3)\n",
    " \n",
    "        tree.write(dotaxml_file, method='xml', encoding='utf-8')  # 更新xml文件\n",
    " \n",
    " \n",
    "# 转换成四点坐标\n",
    "def rotatePoint(xc, yc, xp, yp, theta):\n",
    "    xoff = xp - xc;\n",
    "    yoff = yp - yc;\n",
    "    cosTheta = math.cos(theta)\n",
    "    sinTheta = math.sin(theta)\n",
    "    pResx = cosTheta * xoff + sinTheta * yoff\n",
    "    pResy = - sinTheta * xoff + cosTheta * yoff\n",
    "    return str(int(xc + pResx)), str(int(yc + pResy))\n",
    " \n",
    " \n",
    "def totxt(xml_path, out_path):\n",
    "    # 想要生成的txt文件保存的路径，这里可以自己修改\n",
    " \n",
    "    files = os.listdir(xml_path)\n",
    "    i = 0\n",
    "    for file in files:\n",
    " \n",
    "        tree = ET.parse(xml_path + os.sep + file)\n",
    "        root = tree.getroot()\n",
    " \n",
    "        name = file.split('.')[0]\n",
    " \n",
    "        output = out_path + '/' + name + '.txt'\n",
    "        file = open(output, 'w')\n",
    "        i = i + 1\n",
    "        objs = tree.findall('object')\n",
    "        for obj in objs:\n",
    "            cls = obj.find('name').text\n",
    "            box = obj.find('bndbox')\n",
    "            x0 = int(float(box.find('x0').text))\n",
    "            y0 = int(float(box.find('y0').text))\n",
    "            x1 = int(float(box.find('x1').text))\n",
    "            y1 = int(float(box.find('y1').text))\n",
    "            x2 = int(float(box.find('x2').text))\n",
    "            y2 = int(float(box.find('y2').text))\n",
    "            x3 = int(float(box.find('x3').text))\n",
    "            y3 = int(float(box.find('y3').text))\n",
    "            if x0 < 0:\n",
    "                x0 = 0\n",
    "            if x1 < 0:\n",
    "                x1 = 0\n",
    "            if x2 < 0:\n",
    "                x2 = 0\n",
    "            if x3 < 0:\n",
    "                x3 = 0\n",
    "            if y0 < 0:\n",
    "                y0 = 0\n",
    "            if y1 < 0:\n",
    "                y1 = 0\n",
    "            if y2 < 0:\n",
    "                y2 = 0\n",
    "            if y3 < 0:\n",
    "                y3 = 0\n",
    "            for cls_index, cls_name in enumerate(cls_list):\n",
    "                if cls == cls_name:\n",
    "                    file.write(\"{} {} {} {} {} {} {} {} {} \\n\".format(x0, y0, x1, y1, x2, y2, x3, y3, cls,cls_index))\n",
    "        file.close()\n",
    "        print(output)\n",
    "        print(i)\n",
    " \n",
    " \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "labels2obb/133.txt\n",
      "1\n",
      "labels2obb/144.txt\n",
      "2\n",
      "labels2obb/131.txt\n",
      "3\n",
      "labels2obb/140.txt\n",
      "4\n",
      "labels2obb/118.txt\n",
      "5\n",
      "labels2obb/129.txt\n",
      "6\n",
      "labels2obb/117.txt\n",
      "7\n",
      "labels2obb/108.txt\n",
      "8\n",
      "labels2obb/127.txt\n",
      "9\n",
      "labels2obb/103.txt\n",
      "10\n",
      "labels2obb/124.txt\n",
      "11\n",
      "labels2obb/100.txt\n",
      "12\n",
      "labels2obb/101.txt\n",
      "13\n",
      "labels2obb/134.txt\n",
      "14\n",
      "labels2obb/139.txt\n",
      "15\n",
      "labels2obb/125.txt\n",
      "16\n",
      "labels2obb/106.txt\n",
      "17\n",
      "labels2obb/146.txt\n",
      "18\n",
      "labels2obb/111.txt\n",
      "19\n",
      "labels2obb/147.txt\n",
      "20\n",
      "labels2obb/123.txt\n",
      "21\n",
      "labels2obb/119.txt\n",
      "22\n",
      "labels2obb/120.txt\n",
      "23\n",
      "labels2obb/105.txt\n",
      "24\n",
      "labels2obb/143.txt\n",
      "25\n",
      "labels2obb/1.txt\n",
      "26\n",
      "labels2obb/137.txt\n",
      "27\n",
      "labels2obb/113.txt\n",
      "28\n",
      "labels2obb/138.txt\n",
      "29\n",
      "labels2obb/136.txt\n",
      "30\n",
      "labels2obb/121.txt\n",
      "31\n",
      "labels2obb/114.txt\n",
      "32\n",
      "labels2obb/115.txt\n",
      "33\n",
      "labels2obb/130.txt\n",
      "34\n",
      "labels2obb/112.txt\n",
      "35\n",
      "labels2obb/122.txt\n",
      "36\n",
      "labels2obb/126.txt\n",
      "37\n",
      "labels2obb/132.txt\n",
      "38\n",
      "labels2obb/135.txt\n",
      "39\n",
      "labels2obb/141.txt\n",
      "40\n",
      "labels2obb/148.txt\n",
      "41\n",
      "labels2obb/109.txt\n",
      "42\n",
      "labels2obb/102.txt\n",
      "43\n",
      "labels2obb/110.txt\n",
      "44\n",
      "labels2obb/142.txt\n",
      "45\n",
      "labels2obb/107.txt\n",
      "46\n",
      "labels2obb/145.txt\n",
      "47\n",
      "labels2obb/128.txt\n",
      "48\n",
      "labels2obb/104.txt\n",
      "49\n",
      "labels2obb/149.txt\n",
      "50\n",
      "labels2obb/116.txt\n",
      "51\n"
     ]
    }
   ],
   "source": [
    "if __name__ == '__main__':\n",
    "    # -----**** 第一步：把xml文件统一转换成旋转框的xml文件 ****-----\n",
    "    roxml_path = 'roxml'\n",
    "    dotaxml_path = 'dotaxml'\n",
    "    out_path = 'labels2obb'\n",
    "    filelist = os.listdir(roxml_path)\n",
    "    for file in filelist:\n",
    "        edit_xml(os.path.join(roxml_path, file), os.path.join(dotaxml_path, file))\n",
    " \n",
    "    # -----**** 第二步：把旋转框xml文件转换成txt格式 ****-----\n",
    "    totxt(dotaxml_path, out_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.0rc1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
